#include "main.h"

/*** ELF and PE structures ***/
typedef struct {
  uint8_t  e_ident[16];
  uint16_t e_type;
  uint16_t e_machine;
  uint32_t e_version;
  uint64_t e_entry;
  uint64_t e_phoff;
  uint64_t e_shoff;
  uint32_t e_flags;
  uint16_t e_ehsize;
  uint16_t e_phentsize;
  uint16_t e_phnum;
  uint16_t e_shentsize;
  uint16_t e_shnum;
  uint16_t e_shstrndx;
} __attribute__((packed)) Elf64_Ehdr;

typedef struct {
  uint32_t sh_name;
  uint32_t sh_type;
  uint64_t sh_flags;
  uint64_t sh_addr;
  uint64_t sh_offset;
  uint64_t sh_size;
  uint32_t sh_link;
  uint32_t sh_info;
  uint64_t sh_addralign;
  uint64_t sh_entsize;
} __attribute__((packed)) Elf64_Shdr;

typedef struct {
  uint32_t magic;
  uint16_t machine;
  uint16_t sections;
  uint32_t timestamp;
  uint32_t sym_table;
  uint32_t numsym;
  uint16_t opt_hdr_size;
} __attribute__((packed)) pe_hdr;

typedef struct {
  char     name[8];
  uint32_t vsiz;
  uint32_t vaddr;
  uint32_t rsiz;
  uint32_t raddr;
  uint32_t reloc;
  uint32_t ln;
  uint16_t nreloc;
  uint16_t nln;
  uint32_t chr;
} __attribute__((packed)) pe_sec;

section_t sections[64];

/**
 * Detect if it's an ELF or PE executable
 */
int elfpe_detect(char *fn)
{
    uint8_t tmp[4096];
    uint64_t size;
    void *f = fileopen(fn, &size);

    memset(tmp, 0, sizeof(tmp));
    if(f) {
        fileread(f, tmp, 4096);
        fileclose(f);
    }
    return (tmp[0] == 0x7f && tmp[1] == 'E' && tmp[2] == 'L' && tmp[3] == 'F') || (tmp[0] == 'M' && tmp[1] == 'Z' &&
        *((uint32_t*)&tmp[0x3c]) < sizeof(tmp) - sizeof(pe_hdr) && !memcmp(tmp + *((uint32_t*)&tmp[0x3c]), "PE", 2));
}

/**
 * Read in data sections into sections[] array
 */
int elfpe_getsections(char *fn, int list)
{
    Elf64_Ehdr elf;
    Elf64_Shdr *s;
    pe_hdr pe;
    pe_sec *sec;
    uint8_t tmp[65536];
    uint64_t size;
    uint32_t i, num = 0, strsize;
    void *f = fileopen(fn, &size);
    char *strs = NULL;

    memset(sections, 0, sizeof(sections));
    if(f) {
        fileread(f, tmp, 4096);
        if(tmp[0] == 0x7f && tmp[1] == 'E' && tmp[2] == 'L' && tmp[3] == 'F') {
            memcpy(&elf, tmp, sizeof(Elf64_Ehdr));
            if(elf.e_shoff < size && elf.e_shstrndx < elf.e_shnum && (elf.e_shstrndx + 1) * elf.e_shentsize < (int)sizeof(tmp)) {
                fileseek(f, elf.e_shoff);
                fileread(f, tmp, sizeof(tmp));
                s = (Elf64_Shdr *)(tmp + elf.e_shstrndx * elf.e_shentsize);
                if(s->sh_size && s->sh_offset && (strs = (char*)malloc(s->sh_size))) {
                    strsize = s->sh_size;
                    fileseek(f, s->sh_offset);
                    fileread(f, strs, strsize);
                    for(i = 0; i < elf.e_shnum && (i + 1) * elf.e_shentsize < (int)sizeof(tmp) &&
                        num + 1 < (int)(sizeof(sections)/sizeof(sections[0])); i++) {
                        s = (Elf64_Shdr *)(tmp + i * elf.e_shentsize);
                        if(s->sh_offset && s->sh_size > 64 && s->sh_type == 1 && !(s->sh_flags & ~3) && s->sh_name < strsize &&
                            strs[s->sh_name] && memcmp(strs + s->sh_name, ".eh_", 4) && memcmp(strs + s->sh_name, ".gcc_", 5)) {
                            if(list) debug("section %2u: %8lx %8lx %s", i, s->sh_offset, s->sh_size, strs + s->sh_name);
                            memcpy(&sections[num].name, strs + s->sh_name, 8);
                            sections[num].offs = s->sh_offset;
                            sections[num].size = s->sh_size;
                            num++;
                        }
                    }
                    free(strs);
                } else {
                    engine_features |= EF_NOMEM;
                    debug("unable to allocate memory for ELF section strings (%lu bytes)", s->sh_size);
                }
            } else debug("invalid ELF, no section table '%s'", fn);
        } else
        if(tmp[0] == 'M' && tmp[1] == 'Z' &&
        *((uint32_t*)&tmp[0x3c]) < sizeof(tmp) - sizeof(pe_hdr) && !memcmp(tmp + *((uint32_t*)&tmp[0x3c]), "PE", 2)) {
            memcpy(&pe, tmp + *((uint32_t*)&tmp[0x3c]), sizeof(pe_hdr));
            fileseek(f, *((uint32_t*)&tmp[0x3c]) + pe.opt_hdr_size + 24);
            fileread(f, tmp, sizeof(tmp));
            sec = (pe_sec*)tmp;
            for(i = 0, sec = (pe_sec*)tmp; i < pe.sections && (uint8_t*)sec < tmp + sizeof(tmp) - sizeof(pe_sec) &&
                num + 1 < (int)(sizeof(sections)/sizeof(sections[0])); i++, sec++)
                if(sec->raddr && sec->rsiz > 64 && (sec->chr & 0x40) && !(sec->chr & 0x2000000) && sec->name[0]) {
                    if(list) debug("section %2u: %8x %8x %s", i, sec->raddr, sec->rsiz, sec->name);
                    memcpy(&sections[num].name, sec->name, 8);
                    sections[num].offs = sec->raddr;
                    sections[num].size = sec->rsiz;
                    num++;
                }
        }
        fileclose(f);
    }
    return num;
}
